#!/usr/bin/env python3
import math
import cv2
import apriltag
import numpy as np
import cv_bridge
import rclpy
from rclpy.node import Node
from sensor_msgs.msg import Image
from tf2_ros.transform_broadcaster import TransformBroadcaster
from geometry_msgs.msg import TransformStamped
import tf_transformations as tft

camera_matrix = np.array(
    [
        [1572.569610, 0.000000, 607.879891],
        [0.000000, 1570.246204, 379.644093],
        [0.000000, 0.000000, 1.000000],
    ]
)
dist_coeffs = np.array([0.083966, -0.071645, 0.007297, 0.001435])
cam_params = [1572.569610, 1570.246204, 607.879891, 379.644093]


class AprilTagFinder(Node):
    """
    Image->
    CameraInfo->
    ->TF
    ->operated img
    """

    def __init__(self):
        super().__init__("april_tag_finder")
        self.publisher = self.create_publisher(Image, "image_processed", 1)
        self.subscription = self.create_subscription(
            Image, "camera_image", self.image_callback, 1
        )
        self.tag_detector = apriltag.Detector(
            apriltag.DetectorOptions(families="tag36h11")
        )
        self.tf_broadcaster = TransformBroadcaster(self)

    def image_callback(self, msg):
        cv_image = cv_bridge.CvBridge().imgmsg_to_cv2(msg, desired_encoding="bgr8")
        gray = cv2.cvtColor(cv_image, cv2.COLOR_BGR2GRAY)
        tags = self.tag_detector.detect(gray)
        for tag in tags:
            M, e1, e2 = self.tag_detector.detection_pose(tag, cam_params, 0.2)
            transpose = M[0:3, 3]
            # self.get_logger().debug(
            #    f"x:{transpose[0]}, y:{transpose[1]}, z:{transpose[2]}"
            # )
            t = TransformStamped()
            t.header.stamp = self.get_clock().now().to_msg()
            t.header.frame_id = "camera"
            t.child_frame_id = "tag"
            t.transform.translation.x = -transpose[0]
            t.transform.translation.y = -transpose[1]
            t.transform.translation.z = transpose[2]
            q = tft.quaternion_from_matrix(M.T)
            t.transform.rotation.x = q[0]
            t.transform.rotation.y = q[1]
            t.transform.rotation.z = q[2]
            t.transform.rotation.w = q[3]
            self.tf_broadcaster.sendTransform(t)

            for i in range(4):
                cv2.circle(
                    cv_image, tuple(tag.corners[i].astype(int)), 4, (255, 0, 0), 2
                )
            cv2.circle(cv_image, tuple(tag.center.astype(int)), 4, (2, 180, 200), 4)
        self.publisher.publish(cv_bridge.CvBridge().cv2_to_imgmsg(cv_image, "bgr8"))


def main(args=None):
    rclpy.init(args=args)
    node = AprilTagFinder()
    node.get_logger().info("AprilTag Finder node started")
    rclpy.spin(node)
    rclpy.shutdown()
